package com.soyotec.algorithm.util;

import com.jmatio.io.MatFileHeader;
import com.jmatio.io.MatFileReader;
import com.jmatio.types.MLArray;
import com.jmatio.types.MLDouble;
import com.jmatio.types.MLSingle;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class MatUtil {

    public static Map<String,double[][]> readMatFile(String path){
        return readMatFile(new File(path));
    }

    public static Map<String,double[][]> readMatFile(File file){
        Map<String,double[][]> result = new HashMap<>();
        MatFileReader matFileReader = null;
        try {
            matFileReader = new MatFileReader(file);
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
        // 读取MAT Header信息
        MatFileHeader matFileHeader = matFileReader.getMatFileHeader();
//        System.out.println("*****mat header*****");
//        System.out.println(matFileHeader.toString());
        // 读取矩阵信息,通过矩阵名来获取矩阵信息
//        System.out.println("*****mat body*****");
        Map<String, MLArray> content = matFileReader.getContent();
        for(Map.Entry<String,MLArray> entry : content.entrySet()){
            String key = entry.getKey();
            MLArray mlArray = entry.getValue();
            int[] dimensions1 = mlArray.getDimensions();
            int length = dimensions1.length;
            // 转换为java数组
            if(length==2){
                double[][] transferMLArray2 = transferMLArray2(mlArray);
                result.put(key,transferMLArray2);
            }
        }
        return result;
    }


    /**
     *
     * @param mlArray   把mat文件转换为二维的java的二维
     * @return
     */
    private static double[][] transferMLArray2(MLArray mlArray) {
        int[] dimensions = mlArray.getDimensions();
        double[][] result = null;
        if (mlArray.isSingle()) {
            System.out.println("mlArray.isSingle()");
            MLSingle mlSingle = (MLSingle) mlArray;
            result = new double[dimensions[0]][dimensions[1]];

            for (int i = 0; i < dimensions[0]; i++) {
                for (int j = 0; j < dimensions[1]; j++) {
                    // 列优先，注意转换
                    result[i][j] = mlSingle.get(i,  j);
                }
            }
        } else if (mlArray.isDouble()) {
            System.out.println("mlArray.isDouble()");
            MLDouble mlSingle = (MLDouble) mlArray;
            result = new double[dimensions[0]][dimensions[1]];

            for (int i = 0; i < dimensions[0]; i++) {
                for (int j = 0; j < dimensions[1]; j++) {
                    // 列优先，注意转换
                    result[i][j] = mlSingle.get(i,  j);
                }
            }
        }
        return result;
    }

    /**
     *
     * @param mlArray   把mat文件转换为三维的java的三维
     * @return
     */
    private static double[][][] transferMLArray3(MLArray mlArray) {
        int[] dimensions = mlArray.getDimensions();
        double[][][] result = null;
        if (mlArray.isSingle()) {
            System.out.println("mlArray.isSingle()");
            MLSingle mlSingle = (MLSingle) mlArray;
            result = new double[dimensions[0]][dimensions[1]][dimensions[2]];

            for (int i = 0; i < dimensions[0]; i++) {
                for (int j = 0; j < dimensions[1]; j++) {
                    for (int k = 0; k < dimensions[2]; k++) {
                        // 列优先，注意转换
                        result[i][j][k] = mlSingle.get(i, k * dimensions[1] + j);
                    }
                }
            }
        } else if (mlArray.isDouble()) {
            System.out.println("mlArray.isDouble()");
            MLDouble mlSingle = (MLDouble) mlArray;
            result = new double[dimensions[0]][dimensions[1]][dimensions[2]];

            for (int i = 0; i < dimensions[0]; i++) {
                for (int j = 0; j < dimensions[1]; j++) {
                    for (int k = 0; k < dimensions[2]; k++) {
                        // 列优先，注意转换
                        result[i][j][k] = mlSingle.get(i, k * dimensions[1] + j);
                    }
                }
            }
        }

        return result;
    }
}
